import util
import numpy as np
import Classifier
from sklearn.metrics import (
    accuracy_score,
    precision_score,
    recall_score,
    f1_score,
    classification_report,
)
import joblib,os,time

def trainRF(train_features:np.ndarray,train_labels:np.ndarray,save_path:str):
    classifier=Classifier.BuildClassifier.RF(train_features,train_labels)
    joblib.dump(classifier, save_path)

def trainSVM(train_features:np.ndarray,train_labels:np.ndarray,save_path:str):
        # "linear","rbf","poly"
    prefix=os.path.splitext(save_path)[0]

    linear_classifier=Classifier.BuildClassifier.SVM(train_features,train_labels,"linear")
    joblib.dump(linear_classifier, f"{prefix}_linear.joblib")

    rbf_classifier=Classifier.BuildClassifier.SVM(train_features,train_labels,"rbf")
    joblib.dump(rbf_classifier, f"{prefix}_rbf.joblib")

    poly_classifier=Classifier.BuildClassifier.SVM(train_features,train_labels,"poly")
    joblib.dump(poly_classifier, f"{prefix}_poly.joblib")

def trainPCAKNN(train_features:np.ndarray,train_labels:np.ndarray,save_path:str):
    classifier=Classifier.BuildClassifier.PCAKNN(train_features,train_labels)
    joblib.dump(classifier, save_path)

def testClassifier(classifier,test_features:np.ndarray,test_labels:np.ndarray):
    res= classifier.predict(test_features)
    accuracy = accuracy_score(test_labels, res)
    precision = precision_score(test_labels, res, average="weighted")
    recall = recall_score(test_labels, res, average="weighted")
    f1 = f1_score(test_labels, res, average="weighted")
    report= classification_report(test_labels, res)
    return accuracy,precision,recall,f1,report

def testClassifier(classifier,test_features:np.ndarray,test_labels:np.ndarray):
    # start = time.perf_counter()
    res= classifier.predict(test_features)
    # print(time.perf_counter()-start)
    accuracy = accuracy_score(test_labels, res)
    precision = precision_score(test_labels, res, average="weighted")
    recall = recall_score(test_labels, res, average="weighted")
    f1 = f1_score(test_labels, res, average="weighted")
    report= classification_report(test_labels, res)
    return accuracy,precision,recall,f1,report,res

def train():
    train_features,train_labels=util.load_np_data("train.npz")
    # trainRF(train_features,train_labels,"model/traditional/rf.joblib")
    # trainSVM(train_features,train_labels,"model/traditional/svm.joblib")
    trainPCAKNN(train_features,train_labels,"model/traditional/pcaknn.joblib")

def test():
    test_features,test_labels=util.load_np_data("test.npz")
    # classifier=joblib.load("/home/tuchunxu/workspace/cv_project/vegetable/model/traditional/rf.joblib")
    # classifier=joblib.load("model/traditional/svm_linear.joblib")
    # classifier=joblib.load("model/traditional/svm_poly.joblib")
    # classifier=joblib.load("model/traditional/svm_rbf.joblib")
    # classifier=joblib.load("model/traditional/pcaknn.joblib")
    classifier=joblib.load("pcarf.joblib")
    # print(classifier["pca"].n_components_)
    accuracy,precision,recall,f1,report,prediction=testClassifier(classifier,test_features,test_labels)
    print(f"{accuracy}\n{precision}\n{recall}\n{f1}\n{report}")
    # np.save("res.npy",prediction)

def trainPCARF():
    train_features,train_labels=util.load_np_data("train.npz")
    classifier=Classifier.BuildClassifier.PCAKRF(train_features,train_labels)
    test_features,test_labels=util.load_np_data("test.npz")
    accuracy,precision,recall,f1,report,prediction=testClassifier(classifier,test_features,test_labels)
    print(f"{accuracy}\n{precision}\n{recall}\n{f1}\n{report}")
    joblib.dump(classifier, "pcarf.joblib")

def main():
    test_features,test_labels=util.load_np_data("test.npz")
    classifier=joblib.load("test.joblib")
    accuracy,precision,recall,f1,report=testClassifier(classifier,test_features,test_labels)
    print(f"{accuracy}\n{precision}\n{recall}\n{f1}\n{report}")
    # res=testRF(train_features,train_labels,test_features)
    # print(classification_report,test_labels,res)

if __name__=="__main__":
    # main()
    # train()
    test()
    # trainPCARF()
